#include "ouster_driver.h"

namespace itage {
namespace modules {

OusterLidarDriver::OusterLidarDriver(const std::string& config_file_str) {
    std::cout << "Config file: " << config_file_str << std::endl;

    // 实际配置解析应在此处实现
    // m_lidarIp = config_parser.getLidarIp();
    // m_localIp = config_parser.getLocalIp();

    m_lastDataTime = std::chrono::steady_clock::now();
    m_workerThread = std::thread(&OusterLidarDriver::workerThread, this);
}

OusterLidarDriver::~OusterLidarDriver() {
    m_running = false;
    if (m_workerThread.joinable()) {
        m_cloudReadyCv.notify_all();
        m_workerThread.join();
    }
}

void OusterLidarDriver::workerThread() {
    while (m_running) {
        try {
            poll();
        } catch (const std::exception& e) {
            std::cerr << "Poll error: " << e.what() << ". Reconnecting..." << std::endl;
            std::this_thread::sleep_for(std::chrono::seconds(1));
            reconnectClient();
        }
    }
}

void OusterLidarDriver::reconnectClient() {
    std::lock_guard<std::mutex> lock(m_clientMutex);
    m_client.reset();

    for (int attempt = 0; attempt < 5 && m_running; ++attempt) {
        try {
            m_client = sensor::init_client(m_lidarIp, m_localIp);
            if (m_client) {
                std::cout << "Successfully connected to lidar at: " << m_lidarIp << std::endl;
                return;
            }
        } catch (...) {
        }

        std::this_thread::sleep_for(std::chrono::seconds(1));
    }

    throw std::runtime_error("Failed to reconnect to lidar");
}

void OusterLidarDriver::poll() {
    // 确保客户端连接
    {
        std::lock_guard<std::mutex> lock(m_clientMutex);
        if (!m_client) {
            reconnectClient();
        }
    }

    // 获取元数据
    auto metadata = sensor::get_metadata(*m_client);
    auto info = sensor::parse_metadata(metadata);

    const size_t w = info.format.columns_per_frame;
    const size_t h = info.format.pixels_per_column;

    auto pf = sensor::get_format(info);
    auto scan_batcher = ScanBatcher(info);
    auto ls = LidarScan(w, h, info.format.udp_profile_lidar);
    auto xyz_lut = ouster::make_xyz_lut(info);

    auto packet_buf = std::make_unique<uint8_t[]>(pf.lidar_packet_size);

    // 主数据采集循环
    while (m_running) {
        try {
            client_state st = poll_client(*m_client);

            if (st & CLIENT_ERROR) {
                throw std::runtime_error("Client error detected");
            }

            if (st & LIDAR_DATA) {
                if (read_lidar_packet(*m_client, packet_buf.get(), pf)) {
                    m_lastDataTime = std::chrono::steady_clock::now();

                    if (scan_batcher(packet_buf.get(), ls)) {
                        auto ts_v = ls.timestamp();
                        auto idx = std::find_if(ts_v.begin(), ts_v.end(), [](uint64_t ts) { return ts != 0; });

                        if (idx != ts_v.end()) {
                            auto scan_ts = std::chrono::nanoseconds{*idx};
                            Cloud tmp_cloud{static_cast<uint32_t>(w), static_cast<uint32_t>(h)};

                            scanToCloud(xyz_lut, scan_ts, ls, tmp_cloud, 0);
                            putPointCloud(std::move(tmp_cloud));
                        }
                    }
                }
            }

            // 超时检测
            if (std::chrono::steady_clock::now() - m_lastDataTime > kTimeout) {
                detectTimeout();
            }

            // 适度休眠减少CPU占用
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        } catch (const std::exception& e) {
            std::cerr << "Processing error: " << e.what() << std::endl;
            reconnectClient();
        }
    }
}

void OusterLidarDriver::putPointCloud(Cloud&& tmp_cloud) {
    {
        std::unique_lock<std::mutex> lock(m_cloudMutex);
        m_cloud = std::move(tmp_cloud);
        m_cloudReady = true;
    }
    m_cloudReadyCv.notify_one();
}

bool OusterLidarDriver::getPointCloud(Cloud& out_cloud) {
    std::unique_lock<std::mutex> lock(m_cloudMutex);

    if (m_cloudReadyCv.wait_for(lock, std::chrono::milliseconds(100), [this] { return m_cloudReady; })) {
        out_cloud = std::move(m_cloud);
        m_cloudReady = false;
        return true;
    }
    return false;
}

void OusterLidarDriver::detectTimeout() {
    std::cerr << "Lidar data timeout detected. Attempting to reconnect..." << std::endl;
    reconnectClient();
    m_lastDataTime = std::chrono::steady_clock::now();
}

template <typename T>
inline ouster::img_t<T> get_or_fill_zero(sensor::ChanField f, const ouster::LidarScan& ls) {
    ouster::img_t<T> result{ls.h, ls.w};
    if (ls.field_type(f)) {
        ouster::impl::visit_field(ls, f, read_and_cast(), result);
    } else {
        result = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>::Zero(ls.h, ls.w);
    }
    return result;
}

void OusterLidarDriver::scanToCloud(const ouster::XYZLut& xyz_lut, ouster::LidarScan::ts_t scan_ts,
                                    const ouster::LidarScan& ls, Cloud& cloud, int return_index) {
    const bool second = (return_index == 1);
    cloud.resize(ls.w * ls.h);

    auto near_ir = get_or_fill_zero<uint16_t>(suitableReturn(sensor::ChanField::NEAR_IR, second), ls);

    auto range = get_or_fill_zero<uint32_t>(suitableReturn(sensor::ChanField::RANGE, second), ls);

    auto signal = get_or_fill_zero<uint32_t>(suitableReturn(sensor::ChanField::SIGNAL, second), ls);

    auto reflectivity = get_or_fill_zero<uint16_t>(suitableReturn(sensor::ChanField::REFLECTIVITY, second), ls);

    auto points = ouster::cartesian(range, xyz_lut);
    auto timestamp = ls.timestamp();

// 并行化处理点云数据
#pragma omp parallel for collapse(2)
    for (int u = 0; u < ls.h; u++) {
        for (int v = 0; v < ls.w; v++) {
            const auto xyz = points.row(u * ls.w + v);
            const auto ts = (std::chrono::nanoseconds(timestamp[v]) - scan_ts).count();

            cloud(v, u) = {{{static_cast<float>(xyz(0)), static_cast<float>(xyz(1)), static_cast<float>(xyz(2)), 1.0f}},
                           static_cast<float>(signal(u, v)),
                           static_cast<uint32_t>(ts),
                           static_cast<uint16_t>(reflectivity(u, v)),
                           static_cast<uint8_t>(u),
                           static_cast<uint16_t>(near_ir(u, v)),
                           static_cast<uint32_t>(range(u, v))};
        }
    }
}

ChanField OusterLidarDriver::suitableReturn(sensor::ChanField input_field, bool second) {
    switch (input_field) {
        case sensor::ChanField::RANGE:
        case sensor::ChanField::RANGE2:
            return second ? sensor::ChanField::RANGE2 : sensor::ChanField::RANGE;
        case sensor::ChanField::SIGNAL:
        case sensor::ChanField::SIGNAL2:
            return second ? sensor::ChanField::SIGNAL2 : sensor::ChanField::SIGNAL;
        case sensor::ChanField::REFLECTIVITY:
        case sensor::ChanField::REFLECTIVITY2:
            return second ? sensor::ChanField::REFLECTIVITY2 : sensor::ChanField::REFLECTIVITY;
        case sensor::ChanField::NEAR_IR:
            return sensor::ChanField::NEAR_IR;
        default:
            throw std::runtime_error("Unsupported channel field");
    }
}

}  // namespace modules
}  // namespace itage
